{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "bbad7a3e",
      "metadata": {},
      "source": [
        "# Discrete Bayes Net Example\n",
        "\n",
        "Discrete Bayes Net example with famous Asia Bayes Network"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "colab_button",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/python/gtsam/examples/DiscreteBayesNetExample.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "license_cell",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\nAtlanta, Georgia 30332-0415\nAll Rights Reserved\n\nAuthors: Frank Dellaert, et al. (see THANKS for the full author list)\n\nSee LICENSE for the license information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "colab_import",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n    import google.colab\n    %pip install --quiet gtsam-develop\nexcept ImportError:\n    pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "96168ba4",
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "from gtsam import (DiscreteBayesNet, DiscreteFactorGraph, DiscreteKeys, \n",
        "                   Ordering)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3a2fe874",
      "metadata": {},
      "source": [
        "## Helper Functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "3411f409",
      "metadata": {},
      "outputs": [],
      "source": [
        "def create_discrete_keys(*args):\n",
        "    \"\"\"Create a DiscreteKeys instance from a variable number of DiscreteKey pairs.\"\"\"\n",
        "    dks = DiscreteKeys()\n",
        "    for key in args:\n",
        "        dks.push_back(key)\n",
        "    return dks"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ece48e35",
      "metadata": {},
      "source": [
        "## Asia Bayes Network Example\n",
        "\n",
        "This example demonstrates the famous Asia Bayes Network using discrete Bayes nets.\n",
        "\n",
        "The Asia network is a classic example in probabilistic reasoning that models\n",
        "relationships between:\n",
        "- Visiting Asia (travel history)\n",
        "- Smoking habits\n",
        "- Diseases: Tuberculosis, Lung Cancer, Bronchitis\n",
        "- Symptoms: Dyspnea (shortness of breath)\n",
        "- Tests: X-Ray results\n",
        "\n",
        "The network shows how these variables are conditionally dependent on each other."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fffbe05a",
      "metadata": {},
      "outputs": [],
      "source": [
        "asia = DiscreteBayesNet()\n",
        "\n",
        "# Define discrete keys for each variable (key_id, num_states)\n",
        "Asia = (0, 2)        # Been to Asia: No=0, Yes=1\n",
        "Smoking = (4, 2)     # Smoking: No=0, Yes=1  \n",
        "Tuberculosis = (3, 2) # Tuberculosis: No=0, Yes=1\n",
        "LungCancer = (6, 2)   # Lung Cancer: No=0, Yes=1\n",
        "Bronchitis = (7, 2)   # Bronchitis: No=0, Yes=1\n",
        "Either = (5, 2)       # Either TB or LC: No=0, Yes=1\n",
        "XRay = (2, 2)         # X-Ray positive: No=0, Yes=1\n",
        "Dyspnea = (1, 2)      # Dyspnea: No=0, Yes=1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "dd877e81",
      "metadata": {},
      "outputs": [],
      "source": [
        "# Add prior probabilities\n",
        "asia.add(Asia, \"99/1\")        # P(Asia) = [0.99, 0.01]\n",
        "asia.add(Smoking, \"50/50\")    # P(Smoking) = [0.5, 0.5]\n",
        "\n",
        "# Add conditional probabilities\n",
        "# P(Tuberculosis | Asia)\n",
        "asia.add(Tuberculosis, create_discrete_keys(Asia), \"99/1 95/5\")\n",
        "\n",
        "# P(LungCancer | Smoking)  \n",
        "asia.add(LungCancer, create_discrete_keys(Smoking), \"99/1 90/10\")\n",
        "\n",
        "# P(Bronchitis | Smoking)\n",
        "asia.add(Bronchitis, create_discrete_keys(Smoking), \"70/30 40/60\")\n",
        "\n",
        "# P(Either | Tuberculosis, LungCancer) - OR gate: Either = TB OR LC\n",
        "# \"F T T T\" means: P(Either=1|TB,LC) = [False, True, True, True]\n",
        "# for combinations (TB=0,LC=0), (TB=0,LC=1), (TB=1,LC=0), (TB=1,LC=1)\n",
        "asia.add(Either, create_discrete_keys(Tuberculosis, LungCancer), \"F T T T\")\n",
        "\n",
        "# P(XRay | Either)\n",
        "asia.add(XRay, create_discrete_keys(Either), \"95/5 2/98\")\n",
        "\n",
        "# P(Dyspnea | Either, Bronchitis)\n",
        "asia.add(Dyspnea, create_discrete_keys(Either, Bronchitis), \"9/1 2/8 3/7 1/9\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "d5529655",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Asia \n",
            "size: 8\n",
            "conditional 0:  P( Asia ):\n",
            " Choice(Asia) \n",
            " 0 Leaf 0.99\n",
            " 1 Leaf 0.01\n",
            "\n",
            "conditional 1:  P( Smoking ):\n",
            " Leaf  0.5\n",
            "\n",
            "conditional 2:  P( Tuberculosis | Asia ):\n",
            " Choice(Tuberculosis) \n",
            " 0 Choice(Asia) \n",
            " 0 0 Leaf 0.99\n",
            " 0 1 Leaf 0.95\n",
            " 1 Choice(Asia) \n",
            " 1 0 Leaf 0.01\n",
            " 1 1 Leaf 0.05\n",
            "\n",
            "conditional 3:  P( LungCancer | Smoking ):\n",
            " Choice(LungCancer) \n",
            " 0 Choice(Smoking) \n",
            " 0 0 Leaf 0.99\n",
            " 0 1 Leaf  0.9\n",
            " 1 Choice(Smoking) \n",
            " 1 0 Leaf 0.01\n",
            " 1 1 Leaf  0.1\n",
            "\n",
            "conditional 4:  P( Bronchitis | Smoking ):\n",
            " Choice(Bronchitis) \n",
            " 0 Choice(Smoking) \n",
            " 0 0 Leaf  0.7\n",
            " 0 1 Leaf  0.4\n",
            " 1 Choice(Smoking) \n",
            " 1 0 Leaf  0.3\n",
            " 1 1 Leaf  0.6\n",
            "\n",
            "conditional 5:  P( Either | Tuberculosis LungCancer ):\n",
            " Choice(LungCancer) \n",
            " 0 Choice(Either) \n",
            " 0 0 Choice(Tuberculosis) \n",
            " 0 0 0 Leaf    1\n",
            " 0 0 1 Leaf    0\n",
            " 0 1 Choice(Tuberculosis) \n",
            " 0 1 0 Leaf    0\n",
            " 0 1 1 Leaf    1\n",
            " 1 Choice(Either) \n",
            " 1 0 Leaf    0\n",
            " 1 1 Leaf    1\n",
            "\n",
            "conditional 6:  P( XRay | Either ):\n",
            " Choice(Either) \n",
            " 0 Choice(XRay) \n",
            " 0 0 Leaf 0.95\n",
            " 0 1 Leaf 0.05\n",
            " 1 Choice(XRay) \n",
            " 1 0 Leaf 0.02\n",
            " 1 1 Leaf 0.98\n",
            "\n",
            "conditional 7:  P( Dyspnea | Either Bronchitis ):\n",
            " Choice(Bronchitis) \n",
            " 0 Choice(Either) \n",
            " 0 0 Choice(Dyspnea) \n",
            " 0 0 0 Leaf  0.9\n",
            " 0 0 1 Leaf  0.1\n",
            " 0 1 Choice(Dyspnea) \n",
            " 0 1 0 Leaf  0.3\n",
            " 0 1 1 Leaf  0.7\n",
            " 1 Choice(Either) \n",
            " 1 0 Choice(Dyspnea) \n",
            " 1 0 0 Leaf  0.2\n",
            " 1 0 1 Leaf  0.8\n",
            " 1 1 Choice(Dyspnea) \n",
            " 1 1 0 Leaf  0.1\n",
            " 1 1 1 Leaf  0.9\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Print the network with pretty variable names\n",
        "pretty_names = [\"Asia\", \"Dyspnea\", \"XRay\", \"Tuberculosis\", \n",
        "               \"Smoking\", \"Either\", \"LungCancer\", \"Bronchitis\"]\n",
        "\n",
        "def formatter(key):\n",
        "    return pretty_names[key]\n",
        "\n",
        "asia.print(\"Asia\", formatter)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5b61c652",
      "metadata": {},
      "source": [
        "### Convert to Factor Graph and solve"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "de7cdf8e",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "mpe: (0, 0) (1, 0) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 0)\n",
            "bayesTree: cliques: 6, variables: 8\n",
            "bayesTree- P( Smoking LungCancer Bronchitis ):\n",
            " Choice(Bronchitis) \n",
            " 0 Choice(LungCancer) \n",
            " 0 0 Choice(Smoking) \n",
            " 0 0 0 Leaf 0.3465\n",
            " 0 0 1 Leaf 0.18\n",
            " 0 1 Choice(Smoking) \n",
            " 0 1 0 Leaf 0.0035\n",
            " 0 1 1 Leaf 0.02\n",
            " 1 Choice(LungCancer) \n",
            " 1 0 Choice(Smoking) \n",
            " 1 0 0 Leaf 0.1485\n",
            " 1 0 1 Leaf 0.27\n",
            " 1 1 Choice(Smoking) \n",
            " 1 1 0 Leaf 0.0015\n",
            " 1 1 1 Leaf 0.03\n",
            "\n",
            "bayesTree| - P( Either | LungCancer Bronchitis ):\n",
            " Choice(LungCancer) \n",
            " 0 Choice(Either) \n",
            " 0 0 Leaf 0.9896\n",
            " 0 1 Leaf 0.0104\n",
            " 1 Choice(Either) \n",
            " 1 0 Leaf    0\n",
            " 1 1 Leaf    1\n",
            "\n",
            "bayesTree| | - P( Tuberculosis | Either LungCancer ):\n",
            " Choice(LungCancer) \n",
            " 0 Choice(Either) \n",
            " 0 0 Choice(Tuberculosis) \n",
            " 0 0 0 Leaf    1\n",
            " 0 0 1 Leaf    0\n",
            " 0 1 Choice(Tuberculosis) \n",
            " 0 1 0 Leaf    0\n",
            " 0 1 1 Leaf    1\n",
            " 1 Choice(Either) \n",
            " 1 0 Leaf    0\n",
            " 1 1 Choice(Tuberculosis) \n",
            " 1 1 0 Leaf 0.9896\n",
            " 1 1 1 Leaf 0.0104\n",
            "\n",
            "bayesTree| | | - P( Asia | Tuberculosis ):\n",
            " Choice(Tuberculosis) \n",
            " 0 Choice(Asia) \n",
            " 0 0 Leaf 0.99040016\n",
            " 0 1 Leaf 0.0095998383\n",
            " 1 Choice(Asia) \n",
            " 1 0 Leaf 0.95192308\n",
            " 1 1 Leaf 0.048076923\n",
            "\n",
            "bayesTree| | - P( XRay | Either ):\n",
            " Choice(Either) \n",
            " 0 Choice(XRay) \n",
            " 0 0 Leaf 0.95\n",
            " 0 1 Leaf 0.05\n",
            " 1 Choice(XRay) \n",
            " 1 0 Leaf 0.02\n",
            " 1 1 Leaf 0.98\n",
            "\n",
            "bayesTree| | - P( Dyspnea | Either Bronchitis ):\n",
            " Choice(Bronchitis) \n",
            " 0 Choice(Either) \n",
            " 0 0 Choice(Dyspnea) \n",
            " 0 0 0 Leaf  0.9\n",
            " 0 0 1 Leaf  0.1\n",
            " 0 1 Choice(Dyspnea) \n",
            " 0 1 0 Leaf  0.3\n",
            " 0 1 1 Leaf  0.7\n",
            " 1 Choice(Either) \n",
            " 1 0 Choice(Dyspnea) \n",
            " 1 0 0 Leaf  0.2\n",
            " 1 0 1 Leaf  0.8\n",
            " 1 1 Choice(Dyspnea) \n",
            " 1 1 0 Leaf  0.1\n",
            " 1 1 1 Leaf  0.9\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Convert to factor graph\n",
        "fg = DiscreteFactorGraph(asia)\n",
        "\n",
        "# Create elimination ordering\n",
        "ordering = Ordering()\n",
        "for i in [0, 1, 2, 3, 4, 5, 6, 7]:\n",
        "    ordering.push_back(i)\n",
        "\n",
        "# Solve for most probable explanation (MPE)\n",
        "mpe = fg.optimize()\n",
        "print(\"mpe:\", end=\"\")\n",
        "for i in range(8):\n",
        "    print(f\" ({i}, {mpe[i]})\", end=\"\")\n",
        "print()\n",
        "\n",
        "# Build a Bayes tree (directed junction tree)\n",
        "bayes_tree = fg.eliminateMultifrontal(ordering)\n",
        "bayes_tree.print(\"bayesTree\", formatter)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "222dd5e0",
      "metadata": {},
      "source": [
        "### Add evidence, solve again"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "8ef171de",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "mpe2: (0, 1) (1, 1) (2, 0) (3, 0) (4, 1) (5, 0) (6, 0) (7, 1)\n"
          ]
        }
      ],
      "source": [
        "# Add evidence: we were in Asia and we have dyspnea\n",
        "fg.add(Asia, \"0 1\")      # Evidence: Asia = 1 (Yes, been to Asia)\n",
        "fg.add(Dyspnea, \"0 1\")   # Evidence: Dyspnea = 1 (Yes, have dyspnea)\n",
        "\n",
        "# Solve again with evidence\n",
        "mpe2 = fg.optimize()\n",
        "print(\"mpe2:\", end=\"\")\n",
        "for i in range(8):\n",
        "    print(f\" ({i}, {mpe2[i]})\", end=\"\")\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "939c949a",
      "metadata": {},
      "source": [
        "### Sample from Posterior Distribution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "a7c41812",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "10 samples:\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 1) (5, 0) (6, 0) (7, 1)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 0)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 1) (5, 0) (6, 0) (7, 0)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 1)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 1)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 0)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 1)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 1)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 1)\n",
            "sample: (0, 1) (1, 1) (2, 0) (3, 0) (4, 1) (5, 0) (6, 0) (7, 1)\n"
          ]
        }
      ],
      "source": [
        "chordal = fg.eliminateSequential(ordering)\n",
        "print(\"\\n10 samples:\")\n",
        "for i in range(10):\n",
        "    sample = chordal.sample()\n",
        "    print(\"sample:\", end=\"\")\n",
        "    for j in range(8):\n",
        "        print(f\" ({j}, {sample[j]})\", end=\"\")\n",
        "    print()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}